from exception import NetworkError
from pygeoip import GeoIP   # import module to determine IPlocation
import requests
import json
url = "http://sre-api.gist.symantec.com/sample/network"

class NetworkCompDB:
	def __init__(self, md5list):
		# handles the modifies Md5 list to suit POST format
		self.modifiedMd5list = self.modifyList(md5list)
		
		# handles sample/network results for the list of md5
		# {u'result': [{u'network': [{u'action': u'Search', u'raw': .....}
		self.req = self._getNetworkReq()
		
		# handles the hosts of MD5 set. 
		# e.g. hostMD5ActionDict = {host:{MD5:action,MD5:action}, host:{MD5:action,MD5:action}}
		self.hostMD5ActionDict = {}
		self.hostListOnly = []
		
		# handles the locations based on the IP address
		# e.g. locationMD5IPDict = {location:{MD5:IP,MD5:IP}, location:{MD5:IP,MD5:IP}}
		self.locationMD5IPDict = {}
		self.ipLocationListOnly = []
		
		# handles the network files used as parameter for network based on the raw network
		# e.g. networkFileListOnly = {file:{MD5:raw,MD5:raw}, file:{MD5:raw,MD5:raw}}
		self.networkFileListOnly = []
		self.networkFileMD5RawDict = {}
		
		# start unpacking the network details of the MD5 list
		self._unpackNetworkDetails()

	# modifies the for POST api format
	# e.g { "md5s" : [ "<md5 1>", "<md5 2>", ... "<md5 n>" ]}
	def modifyList(self, md5list):
		data = {}
		listOfMd5String = []
		for item in md5list:
			listOfMd5String.append(item)
		data["md5s"] = listOfMd5String
		return data
	
	# sends a request for the sample/network using POST RATS api
	# returns req in json format
	def _getNetworkReq(self):
		try:
			req = requests.post(url, data=json.dumps(self.modifiedMd5list))
# 			print json.dumps(req.json(), indent =2)
			return req.json()
		except ValueError:
			raise NetworkError("Produces bad Network Request")
	
	# unpack network details, call for each unpacking methods
	def _unpackNetworkDetails(self):
		for item in self.req['result']:
			nameOfMd5 = item['md5']
			for networkItem in item['network']:
				self._unpackHost(networkItem, nameOfMd5)
				self._unpackIPLocation(networkItem, nameOfMd5)
				self._unpackNetworkFile(networkItem, nameOfMd5)								
		self.hostListOnly = list(set(self.hostListOnly)) 					# make sure no duplicates
		self.ipLocationListOnly = list(set(self.ipLocationListOnly)) 		# make sure no duplicates
		self.networkFileListOnly = list(set(self.networkFileListOnly)) 		# make sure no duplicates
	
	# unpacking host of network
	def _unpackHost(self, networkItem, nameOfMd5):
		md5action = {}
		actionlist = []
		nameOfHost = networkItem['host']
		self.hostListOnly.append(nameOfHost)
		if self.hostMD5ActionDict.has_key(nameOfHost):
			md5action = self.hostMD5ActionDict[nameOfHost]
			if md5action.has_key(nameOfMd5):
				actionlist = md5action[nameOfMd5]
		actionlist.append(networkItem['action'])
		md5action[nameOfMd5] = actionlist
		self.hostMD5ActionDict[nameOfHost] = md5action	# set {network:{md5:actions}}
		self.hostListOnly.append(nameOfHost)			# append network found in the set of MD5s to hostListOnly
		
	# unpacking IPlocation
	def _unpackIPLocation(self, networkItem, nameOfMd5):
		gic = GeoIP('GeoLiteCity.dat') 	# GEO lite database
		md5ip = {}
		if networkItem['as'] is not None:
			ip = networkItem['as']['ip']
			ipLocation = gic.time_zone_by_addr(networkItem['as']['ip'])
			ipList = []
			if self.locationMD5IPDict.has_key(ipLocation):
				md5ip = self.locationMD5IPDict[ipLocation]
				if md5ip.has_key(nameOfMd5):
					ipList = md5ip[nameOfMd5]
			ipList.append(ip)							
			ipList = list(set(ipList))
			md5ip[nameOfMd5] = ipList
			self.locationMD5IPDict[ipLocation] = md5ip
# 						print self.locationMD5IPDict
			self.ipLocationListOnly.append(ipLocation)
	
	# unpacking network file parameters
	# assuming a network file is contained at the end of a string that contains '/'
	# e.g. "www.google.com/example.exe" -> network file is "example.exe"
	def _unpackNetworkFile(self, networkItem, nameOfMd5):
		rawNetwork = networkItem['raw']
		md5raw = {}
		splitNetwork = rawNetwork.split('/')		# split by '/'
		if len(splitNetwork) > 1 :					# file is at the end of the list
			fileParam = splitNetwork.pop()			
			if fileParam and '.' in fileParam:
				rawList = []
				if self.networkFileMD5RawDict.has_key(fileParam):
					md5raw = self.networkFileMD5RawDict[fileParam]
					if md5raw.has_key(nameOfMd5):
						rawList = md5raw[nameOfMd5]
				rawList.append(rawNetwork)
				md5raw[nameOfMd5] = rawList
				self.networkFileMD5RawDict[fileParam] = md5raw
				self.networkFileListOnly.append(fileParam)
	
	# compare the list of host from the two lists of MD5s
	# takes as parameter the list of host from other MD5 list
	# returns list of similarity of hosts
	def compareHost(self, networkList_other):
		hostInListA = []
		hostInListB = []
		for item in self.hostListOnly:
			hostInListA.append(str(item))
		for item in networkList_other:
			hostInListB.append(str(item))
		similarity = list(set(hostInListA) & set(hostInListB))
		return similarity
	
	# compares the list of IP locations between two lists of MD5s
	# takes as parameter the list of location of other MD5 list
	# returns list of similarity of locations
	def compareLocation(self, locationList_other):
		similarity = list(set(self.ipLocationListOnly) & set(locationList_other))
# 		print "similarity : ", similarity
		return similarity
		
	# compares the list of network files between two lists of MD5s
	# takes as parameter the list of network files of other MD5 list
	# returns list of similarity of network files
	def compareNetworkFile(self, networkFileList_other):
		similarity = list(set(self.networkFileListOnly) & set(networkFileList_other))
		return similarity
		
			